## 从数据集中根据标签名称找到对应的图片
import os
import shutil
import xml.etree.ElementTree as ET

# 标签文件夹路径
annotations_folder = '/home/shuai/Downloads/public_dataset/SeaShips(7000)/Annotations'
# 图片文件夹路径
images_folder = '/home/shuai/Downloads/public_dataset/SeaShips(7000)/JPEGImages'
# 目标文件夹路径
output_folder = '/home/shuai/Downloads/public_dataset/SeaShips(7000)/filtered_images'

# 如果目标文件夹不存在，则创建它
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 遍历标签文件夹中的所有.xml文件
for filename in os.listdir(annotations_folder):
    if filename.endswith('.xml'):
        # 解析xml文件
        xml_path = os.path.join(annotations_folder, filename)
        tree = ET.parse(xml_path)
        root = tree.getroot()
        
        # 检查是否有且只有一个分类为"ore carrier"
        objects = root.findall('object')
        if len(objects) == 1 and objects[0].find('name').text == 'fishing boat':
            # 获取图片文件名
            image_filename = root.find('filename').text
            image_path = os.path.join(images_folder, image_filename)
            
            # 如果图片存在，将其复制到目标文件夹
            if os.path.exists(image_path):
                shutil.copy(image_path, os.path.join(output_folder, image_filename))

print("图片已成功提取到目标文件夹。")
